# _*_ coding:utf-8 _*_
# @Time  : 2023.10.24
# @Author: zizlee
from typing import List

import pandas as pd
from fastapi import APIRouter, Depends, Body
from fastapi.encoders import jsonable_encoder
from pydantic import BaseModel
from response import APIResponse
from constant.section_scope import SectionScope
from db_utils import BaseDatabaseConnection, IndexDataBaseConnection
from utils import datetime_utils
from utils.jwt_token import TokenPayload
from utils.module_public import ModulePublic
from utils.validate_util import DateRangeValidator

spread_arbitrage_api = APIRouter()


class ExConfigItem(BaseModel):
    vcode: str
    vname: str
    series_name: str
    direction: str
    db_name: str
    tb_name: str
    multiplier: float
    volume: float
    direction: str


class SpreadArbitrageCalculateItem(DateRangeValidator):
    last_years: int
    config_items: List[ExConfigItem]


class VarietyAuth:
    def user_variety(self, payload: TokenPayload):
        public_util = ModulePublic(module_key=SectionScope.SPREAD_ARBITRAGE_VARIETY_SCOPE, module_sec=SectionScope.SPREAD_ARBITRAGE_KEY)
        variety_list = public_util.access_variety(payload=payload)
        # print(variety_list)
        # 获取品种下的交割月
        sql = """
            SELECT id,vcode,delivery,db_name,tb_name,dat_name FROM datlib_index_mapper
            WHERE flag='spreadArbitrage' AND is_active=1
            ORDER BY sort_num;
        """
        with BaseDatabaseConnection() as conn:
            conn.cursor.execute(sql)
            mapper_data = list(conn.cursor.fetchall())
        for m in mapper_data:
            for v_obj in variety_list:
                if m['vcode'] == v_obj['code']:
                    v_obj.setdefault('delivery', []).append(m)
        return APIResponse.success(data=variety_list)


class SpreadArbitrageCalculate:

    @staticmethod
    def echarts_common_option(name, unit: str = ''):
        # color: ['#5470c6', '#91cc75', '#fac858', '#ee6666', '#73c0de', '#3ba272', '#fc8452', '#9a60b4', '#ea7ccc']
        return {
            'title': {
                'text': name,
                'left': '10%',
                'top': 10,
            },
            'grid': {
                'top': 60,
                'left': '10%',
                'bottom': 53,
            },
            'legend': {
                'type': 'scroll',
                'bottom': 7,
                'left': '10%'
            },
            'xAxis': {
                'type': 'category',
                'data': []
            },
            'yAxis': [
                {
                    'type': 'value',
                    'name': unit,
                    'splitLine': {'show': True},
                    'scale': True
                }
            ],
            'axisPointer': {
                'show': True
            },
            'series': [],
            'toolbox': {
                'show': True,
                'feature': {
                    'magicType': {
                        'type': ['line', 'bar'],
                        'title': {
                            'line': '切换为折线图',
                            'bar': '切换为柱状图',
                            'stack': '切换为堆叠',
                            'tiled': '切换为平铺'
                        }
                    },
                    'restore': {
                        'title': '还原'
                    },
                    'saveAsImage': {
                        'title': '保存为图片'
                    }
                }
            },
            'tooltip': {
                'trigger': 'axis',
                'showContent': True,
                'axisPointer': {
                    'type': 'cross',
                }
            },
        }

    def calculating(self, config_data):
        if not all([config_data.start_date, config_data.end_date]):
            return APIResponse.validate_error(detail='请设置正确的开始和结束日期!')
        # 获取数据
        index_data_list = self._get_database(start_date=config_data.start_date, end_date=config_data.end_date,
                                             config_rows=jsonable_encoder(config_data.config_items))
        print(config_data)
        # 数据的计算处理和生成图形配置
        df = None
        spread_arbitrage_formula = ''
        daily_profit_formula = ''
        column_names = {}
        for index, row in enumerate(index_data_list):
            tmp = pd.DataFrame(row['data'])
            if df is not None and tmp.empty:
                continue
            tmp = tmp[tmp['datavalue'] != 0]
            if tmp.empty:
                continue
            tmp[f'multiplier_{index}'] = row['multiplier']
            tmp[f'volume_{index}'] = row['volume']
            tmp[f'direction_{index}'] = row['direction']
            tmp[f'diff_{index}'] = tmp['datavalue'] - tmp.loc[0, 'datavalue']
            if row['direction'] == 'S':  # 空
                spread_arbitrage_formula += f'+value_{index}'
                daily_profit_formula += f'-diff_{index}*multiplier_{index}*volume_{index}'
            else:
                spread_arbitrage_formula += f'-value_{index}'
                daily_profit_formula += f'+diff_{index}*multiplier_{index}*volume_{index}'
            column_names[f'value_{index}'] = '{}-{}'.format(row['vname'], row['series_name'])
            tmp.rename(columns={'datavalue': f'value_{index}'}, inplace=True)
            # print(tmp)
            del tmp['id']
            if df is None:
                df = tmp
            else:
                df = df.merge(tmp, on='datadate')
        df.eval(f'spread_arbitrage = {spread_arbitrage_formula}', inplace=True)  # 计算价差(空合计-多合计)
        df.dropna(inplace=True)  # 计算结果有问题的删除
        # 生成普通价差图形
        basic_opt = self.echarts_common_option(name='品种价差分析')
        basic_opt['xAxis']['data'] = df['datadate'].to_list()
        series_data = []
        legend_data = []
        for col, name in column_names.items():
            series_data.append({
                'name': name,
                'type': 'line',
                'symbol': 'none',
                'connectNulls': True,
                'data': df[col].to_list()
            })
            legend_data.append(name)
        if len(basic_opt['yAxis']) < 2:
            basic_opt['yAxis'].append({
                'type': 'value',
                'name': '价差',
                'splitLine': {'show': False},
                'scale': True
            })
        series_data.append({
            'name': '价差',
            'type': 'bar',
            'data': df['spread_arbitrage'].to_list(),
            'yAxisIndex': 1
        })

        basic_opt['legend']['data'] = legend_data
        basic_opt['series'] = series_data
        # print(spread_arbitrage_formula)
        # print(df)
        # print(config_data)
        # basic_opt = self.generate_spread_chart_option(data_df=df.copy())

        # 生成价差季节图形
        spread_arbitrage_list = df[['datadate', 'spread_arbitrage']].to_dict(orient='records')
        # 分出年份
        years_data = {}
        for sa in spread_arbitrage_list:
            y = sa['datadate'][:4]
            years_data.setdefault(y, {})[sa['datadate'][-5:]] = round(sa['spread_arbitrage'], 2)
        season_opt = self.echarts_common_option(name='价差季节分析')
        x_data_season = datetime_utils.generate_dates_of_year()
        legend_data_season = []
        series_data_season = []
        for y, y_data in years_data.items():
            legend_data_season.append(y)
            series_data_season.append({
                'name': y,
                'type': 'line',
                'symbol': 'none',
                'connectNulls': True,
                'data': [y_data.get(x) for x in x_data_season]
            })
        season_opt['legend']['data'] = legend_data_season
        season_opt['series'] = series_data_season
        season_opt['xAxis']['data'] = x_data_season

        # 当日收益试算
        df.eval(f'profit = {daily_profit_formula}', inplace=True)  # 试算每日收益：第一日按条件操作，至当日的收益试算
        # print(df)
        # print(daily_profit_formula)
        profit_opt = self.echarts_common_option(name='当日收益试算分析')
        profit_opt['xAxis']['data'] = df['datadate'].to_list()
        profit_opt['series'] = [{
            'name': '收益',
            'type': 'bar',
            'symbol': 'none',
            'connectNulls': True,
            'data': [round(v, 2) for v in df['profit'].to_list()]
        }]
        return APIResponse.success(data={'basic_opt': basic_opt, 'season_opt': season_opt, 'profit_opt': profit_opt})

    def _get_database(self, start_date, end_date, config_rows):
        # 整理共同的库的指标
        db_indexes = {}
        for iid in config_rows:
            iid['fs'] = start_date
            iid['fe'] = end_date
            db_indexes.setdefault(iid['db_name'], []).append(iid)
        index_data_list = []
        for db_name, index_list in db_indexes.items():
            with IndexDataBaseConnection(db_name=db_name) as conn:
                for o_idx in index_list:
                    o_data = conn.query_with_interval(tb_name=o_idx['tb_name'],
                                                      start_date=o_idx['fs'], end_date=o_idx['fe'],
                                                      ascending=True)
                    index_data_list.append({**o_idx, 'data': o_data['rows']})
        return index_data_list


@spread_arbitrage_api.get('/variety/')  # 获取有权限的品种列表
def user_variety(payload: TokenPayload = Depends(TokenPayload)):
    return VarietyAuth().user_variety(payload=payload)


@spread_arbitrage_api.post('/calculating/')  # 计算价差套利
def calculate_spread_arbitrage(config: SpreadArbitrageCalculateItem = Body(...)):
    return SpreadArbitrageCalculate().calculating(config_data=config)

